# %%
# 用三阶多项式拟合y = sin(x)
# 欧几里德距离最小化来进行梯度下降训练
# %%
# 用numpy来操作手动实现网络前后传播
import numpy as np
import math

# 创建随机的输入和输出数据
# 从-pi到pi，两端都要，线性均匀取2000个点
x = np.linspace(-math.pi, math.pi, 2000)
y = np.sin(x)

# 随机初始化权重,有a,b,c,d四个参数
# 正态分布初始化
a = np.random.randn()
b = np.random.randn()
c = np.random.randn()
d = np.random.randn()

# 设置学习率
learning_rate = 1e-6
# 训练2000轮
for t in range(2000):
    # 前向传播，计算y
    # y = a + b*x + c*x^2 + d*x^3
    y_pred = a + b * x + c * x ** 2 + d * x ** 3
    # 计算然后打印出loss,这里采用的是MSE,并求和
    # loss = (y_pred-y)^2
    loss = np.square(y_pred - y).sum()
    # 因为从0开始那么99刚好就是100轮了
    if t % 100 == 99:
        print(t, loss)
    # 反向传播然后计算参数a,b,c,d对loss的梯度
    # 因为loss = (y_pred-y)^2，所以grad_y_pred = 2.0*(y_pred-y)
    grad_y_pred = 2.0 * (y_pred - y)
    # 而y_pred = a + b*x + c*x^2 + d*x^3
    # 所以grad_x对loss就是复合函数求导了
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()
    # 求导之后，来更新参数吧
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d
# 最后打印出结果
print(f'Result: y = {a} + {b} x + {c} x^2 + {d} x^3')
# %%
# 这里用tensor来替换ndarray重复上面的例子
# 并搬到GPU上运算，假装有GPU
# 其实没有太多的区别甚至基本上可以用torch替换np，也许就是为什么pytorch受欢迎吧
import torch
import math

# 设定在CPU/GPU上运算，在定义每个变量的时候都加一个
# dtype=dtype, device=device 自然就在对应的设备上了
dtype = torch.float
device = torch.device("cpu")
# 取消注释就在GPU0上运算了
# device = torch.device("cuda:0")

# 创建随机的输入和输出数据
# 从-pi到pi，两端都要，线性均匀取2000个点
x = torch.linspace(-math.pi, math.pi, 2000,
                   dtype=dtype, device=device)
y = torch.sin(x)

# 随机初始化权重,有a,b,c,d四个参数
# 正态分布初始化
# 只是说这里有点区别，要先填形状，标量没有形状的话，也要拿个()占位
a = torch.randn((), dtype=dtype, device=device)
b = torch.randn((), dtype=dtype, device=device)
c = torch.randn((), dtype=dtype, device=device)
d = torch.randn((), dtype=dtype, device=device)

# 设置学习率
learning_rate = 1e-6
# 训练2000轮
for t in range(2000):
    # 前向传播，计算y
    # y = a + b*x + c*x^2 + d*x^3
    y_pred = a + b * x + c * x ** 2 + d * x ** 3
    # 计算然后打印出loss,这里采用的是MSE,并求和
    # loss = (y_pred-y)^2
    loss = torch.square(y_pred - y).sum()
    # 因为从0开始那么99刚好就是100轮了
    if t % 100 == 99:
        print(t, loss)
    # 反向传播然后计算参数a,b,c,d对loss的梯度
    # 因为loss = (y_pred-y)^2，所以grad_y_pred = 2.0*(y_pred-y)
    grad_y_pred = 2.0 * (y_pred - y)
    # 而y_pred = a + b*x + c*x^2 + d*x^3
    # 所以grad_x对loss就是复合函数求导了
    grad_a = grad_y_pred.sum()
    grad_b = (grad_y_pred * x).sum()
    grad_c = (grad_y_pred * x ** 2).sum()
    grad_d = (grad_y_pred * x ** 3).sum()
    # 求导之后，来更新参数吧
    a -= learning_rate * grad_a
    b -= learning_rate * grad_b
    c -= learning_rate * grad_c
    d -= learning_rate * grad_d
# 最后打印出结果
print(f'Result: y = {a} + {b} x + {c} x^2 + {d} x^3')
# %%
# 张量和Autograd
# 利用Autograd来实现反向传播
import torch
import math

dtype = torch.float
device = torch.device("cpu")

# 创建数据集
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)

# 如果是自己用张量定义参数的话那么requires_grad = True,不然不会求导
a = torch.randn((), device=device, dtype=dtype, requires_grad=True)
b = torch.randn((), device=device, dtype=dtype, requires_grad=True)
c = torch.randn((), device=device, dtype=dtype, requires_grad=True)
d = torch.randn((), device=device, dtype=dtype, requires_grad=True)

# 打印下数据集的形状
print('x', x.size())
print('y', y.size())

learning_rate = 1e-6
for t in range(2000):
    # 其实前向传播就是在定义网络结构
    y_pred = a + b * x + c * x ** 2 + d * x ** 3
    # 计算loss
    # 用tensor的operations来计算
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        # loos的shape是(1,)的tensor，要用loss.item()把值拿出来成为标量
        print(t, loss.item())
    # 调用backend()就会计算loss对所有requires_grde=True变量的梯度，会保存起来
    loss.backward()
    # 我们进行跟新参数运算时，就不用再构建计算图了
    # 为什么要加 torch.no_grad()？
    # 如果不加，torch会默认把更新参数的运算也加入到计算图了
    # 下次backward()连更新权重的运算也会被计算梯度，这是错误的。
    with torch.no_grad():
        a -= learning_rate * a.grad
        b -= learning_rate * b.grad
        c -= learning_rate * c.grad
        d -= learning_rate * d.grad

        # 清除累计在参数上的梯度
        a.grad = None
        b.grad = None
        c.grad = None
        d.grad = None
print(f'Result: y = {a.item()} + {b.item()} x + {c.item()} x^2 + {d.item()} x^3')

# %%
"""自定义Autograd函数，不是很懂 2021年7月24日"""
# TODO 我不明白怎么自定义Autograd
# 自定义Autograd函数
# 模型改成用y = a + b P[3](c + dx)来拟合
# P[3](x) = 1/2 (5x ^ 3 - 3x)
import torch
import math


class LegendrePolynomial3(torch.autograd.Function):
    # 我们能够实现自己的autograd函数，通过定义torch.autograd.Function
    # 的子类，但是要用tensor的operator实现forward和backward

    # 静态方法，可以不实例化就调用
    @staticmethod
    def forward(ctx, input):
        # 在前向传播中，接受input的tensor输出output的tensor
        # ctx  context object 上下文对象，它作用是：
        # 来储存信息用于backward的计算
        # ctx.save_for_backward方法可以用来缓存任意的对象
        ctx.save_for_backward(input)
        return 0.5 * (5 * input ** 3 - 3 * input)

    @staticmethod
    def backward(ctx, grad_output):
        input, _ = ctx.saved_tensors
        return grad_output * 1.5 * (5 * input ** 2 - 1)


dtype = torch.float
device = torch.device('cpu')
# 创建数据集
x = torch.linspace(-math.pi, math.pi, 2000, device=device, dtype=dtype)
y = torch.sin(x)
# 创建参数a,b,c,d  y = a + b * P3(c + d * x)
# 把这些参数的初始化的离正确答案很近，确保收敛性
# 为什么用torch.full()，目的是为了离正确答案近一点保证收敛性
a = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
b = torch.full((), -1.0, device=device, dtype=dtype, requires_grad=True)
c = torch.full((), 0.0, device=device, dtype=dtype, requires_grad=True)
d = torch.full((), 0.3, device=device, dtype=dtype, requires_grad=True)
learning_rate = 5e-6
for t in range(2000):
    # 为了使用我们的函数，我们用P3来命名
    P3 = LegendrePolynomial3.apply
    # 前向传播计算y_pred,用自己的函数来计算P3
    y_pred = a + b * P3(c + d * x)
    # 计算并打印loss
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())

    # 用自动求导反向传播
    loss.backward()
    # 用梯度下降来更新参数
    with torch.no_grad():
        a -= learning_rate * a.grad
        b -= learning_rate * b.grad
        c -= learning_rate * c.grad
        d -= learning_rate * d.grad

        # 清除pytorch的累计梯度
        a.grad = None
        b.grad = None
        c.grad = None
        d.grad = None

print(f'Result: y = {a.item()} + {b.item()} * P3({c.item()} + {d.item()} x)')
